#include <iostream>
#include <sstream>
#include <fstream>
#include <vector>
#include <string>
#include <cstring>
#include <cstdlib>

#include <sys/mman.h>
#include <sys/time.h>
#include <dirent.h>
#include <unistd.h>
#include <errno.h>
#include <pthread.h>

using namespace std;

void * testThread(void * arg);
void createDettachThread(pthread_t & tid, void* (* fun) (void *), void * arg);
bool getFileNamesInDir(const string& strDir, vector<string>& vecFileName);
double timeDiff(struct timeval & startTime, struct timeval & endTime);
size_t getFileSize(string filename);

typedef struct targ {
    int id;
    vector<string> * filenames;
    string diskname;
} testArg;

// GLOBAL VARIBLES
int acctime;
size_t blockSize;

int main(int argc, char **argv)
{

    int threadNum = 0;
    int memoryBlockGB = 0;
    string diskname[] = {"disk0", "disk1", "disk2", 
                         "disk3", "disk4", "disk5"};
    vector<string> filenames[6];
    blockSize = 0;
    acctime = 0;

    if (argc != 5) {
        cout << "Usage : randtest [threadnum] [acctime] [blocksize] [memoryblockupGB]" << endl;
        exit(-1);
    }
    threadNum = atoi(argv[1]);
    acctime   = atoi(argv[2]);
    blockSize = atoll(argv[3]);
    memoryBlockGB = atoi(argv[4]);

    // consume memory
    size_t totalSize = 1024 * 1024 * 1024 ;
    char * mStartAddr = (char *)malloc( memoryBlockGB * totalSize );
    int ret = mlock((void *)mStartAddr, memoryBlockGB * totalSize);

    pthread_t tid;
    for(int i=0; i < 6; i++) {
        if (!getFileNamesInDir("/" + diskname[i], filenames[i])) {
            cout << "Error : get dir info fail." << endl;
            exit(-1);
        }

        if (filenames[i].size() == 0) {
            cout << "Error\t: There are no test file fit the blocksize in DIR [ " 
                 << diskname[i] << " ]."  << endl
                 <<  "\tNeed add test file first ." << endl;
            exit(-1);
        }

        for(int j=0; j < threadNum; j++) {
            testArg * threadarg = new targ();
            threadarg->id = j;
            threadarg->filenames = &filenames[i];
            threadarg->diskname = diskname[i];
            createDettachThread(tid, testThread, (void*) threadarg);
        }
    }

    sleep(acctime + 10);

    munlockall();

    return 0;
}

void set_pthread_dettach(pthread_attr_t & attr)
{
    if ( pthread_attr_init( &attr ) != 0 ) {
        cout << "\tError : pthread_attr_init fail " << strerror(errno) << endl;
        return;
    }

    if ( pthread_attr_setdetachstate(&attr, PTHREAD_CREATE_DETACHED) != 0) {
        cout << "\tError : pthread_attr_setdetachstate fail " << strerror(errno) << endl;
        return;
    }
}

void createDettachThread(pthread_t & tid, void* (* fun) (void *), void * arg)
{
    pthread_attr_t attr;
    set_pthread_dettach( attr );
    pthread_create(&tid, &attr, fun, arg);
}

void * testThread(void * arg)
{
    testArg * testarg = (testArg *)arg;
    int threadid      = testarg->id;
    string diskname   = testarg->diskname;
    vector<string> & filenames = *(testarg->filenames);
    ifstream file;

    int i, filenum = filenames.size();
    struct timeval startTime, endTime, lastTime;
    size_t filelength = 0, totalread = 0, secread = 0;
    char buf[4096];

    // record speed
    ofstream recordfile;
    stringstream ss;
    ss << "t_" << diskname << "_" << threadid << ".dat";
    recordfile.open(ss.str().c_str());

    gettimeofday(&startTime, NULL);
    lastTime.tv_sec = startTime.tv_sec;
    lastTime.tv_usec = startTime.tv_usec;

    while(true) {
        i = rand() % filenum;
        string name = "/" + diskname + "/" + filenames[i];
        file.open( name.c_str() );
        if (file.is_open()) {

            file.seekg(0, ios::end);
            filelength = file.tellg();
            file.seekg(0, ios::beg);

            totalread += filelength;
            secread   += filelength;

            while(filelength > 0) {
                if (filelength >= 4096) {
                    file.read(buf, 4096);
                    filelength -= 4096;
                } else {
                    file.read(buf, filelength);
                    filelength = 0;
                }
            }
            file.close();
        }

        gettimeofday(&endTime, NULL);
        double diffsec = timeDiff(lastTime, endTime);
        double diffall = timeDiff(startTime, endTime);

        if ( diffsec >= 1.0) {
            recordfile << diffall << " " << (secread / 1024.0 / 1024.0) / diffsec << endl;
            secread = 0;
            lastTime.tv_sec = endTime.tv_sec;
            lastTime.tv_usec = endTime.tv_usec;
        }

        if ( diffall >= acctime ) {
            recordfile << diffall << " " << (totalread / 1024.0 / 1024.0) / diffall << endl;
            recordfile.close();
            break;
        } 
    }

    return (void *)0;
}

double timeDiff(struct timeval & startTime, struct timeval & endTime)
{
    return ((endTime.tv_sec - startTime.tv_sec) * 1000000 + 
            (endTime.tv_usec - startTime.tv_usec)) / 1000000.0;
}

bool getFileNamesInDir(const string & strDir,
                       vector<string>& vecFileName)
{
    DIR* dir = NULL;
    struct dirent entry;
    struct dirent* entryPtr = NULL;

    dir = opendir(strDir.c_str());
    if (NULL == dir)
    {
        cout << strerror(errno) << ", strDir : " << strDir << endl;
        return false;
    }

    readdir_r(dir, &entry, &entryPtr);
    while (entryPtr != NULL) 
    {
        if (entry.d_type == DT_REG) {
            string strFileName = entry.d_name;
            if ("." == strFileName || ".." == strFileName)
            {
            }
            else
            {
                if (getFileSize( "/" + strDir + "/" + strFileName ) == blockSize) {
                    vecFileName.push_back(strFileName);
                }
            }
        }
        readdir_r(dir, &entry, &entryPtr);
    }

    return true;
}

size_t getFileSize(string filename)
{
    size_t size;
    ifstream f;
    f.open(filename.c_str());

    if (f.is_open()) {
        f.seekg(0, ios::end);
        size = f.tellg();
        f.seekg(0, ios::beg);
        f.close();
    } else {
        f.close();
        size = 0;
    }

    return size;
}
